import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import random

# 参数设置
grid_size = 50
num_people = 100
exit1 = (0, 0)
exit2 = (grid_size-1, grid_size-1)
obstacle_start = (10, 10)
obstacle_end = (20, 20)

# 初始化网格
grid = np.zeros((grid_size, grid_size), dtype=int)
# 设置出口
grid[exit1] = 3
grid[exit2] = 3
# 设置障碍物
for i in range(obstacle_start[0], obstacle_end[0]):
    for j in range(obstacle_start[1], obstacle_end[1]):
        grid[i][j] = 2

# 初始化人群
people = []
for _ in range(num_people):
    while True:
        row = random.randint(0, grid_size-1)
        col = random.randint(0, grid_size-1)
        if grid[row][col] == 0:
            grid[row][col] = 1
            people.append({'position': (row, col), 'history': [(row, col)]})
            break

# 生成势场
def generate_field(grid, exit_positions):
    from collections import deque
    rows, cols = grid.shape
    visited = np.full((rows, cols), False)
    distance = np.full((rows, cols), np.inf)
    
    for exit_pos in exit_positions:
        queue = deque()
        queue.append(exit_pos)
        visited[exit_pos] = True
        distance[exit_pos] = 0
        
        while queue:
            current = queue.popleft()
            for dr, dc in [(-1,0),(1,0),(0,-1),(0,1)]:
                nr, nc = current[0]+dr, current[1]+dc
                if 0 <= nr < rows and 0 <= nc < cols:
                    if not visited[nr][nc] and grid[nr][nc] != 2:
                        visited[nr][nc] = True
                        distance[nr][nc] = distance[current] + 1
                        queue.append((nr, nc))
    
    return distance

field = generate_field(grid, [exit1, exit2])

# 更新规则
def update_people(grid, field, people):
    new_people = []
    for p in people:
        row, col = p['position']
        min_distance = np.inf
        best_move = None
        
        for dr, dc in [(-1,0),(1,0),(0,-1),(0,1)]:
            nr, nc = row+dr, col+dc
            if 0 <= nr < grid_size and 0 <= nc < grid_size:
                if grid[nr][nc] == 0 and field[nr][nc] < min_distance:
                    min_distance = field[nr][nc]
                    best_move = (nr, nc)
        
        if best_move:
            # 更新位置
            grid[row][col] = 0
            grid[best_move] = 1
            p['position'] = best_move
            p['history'].append(best_move)
            new_people.append(p)
        else:
            # 无可行移动，保持原地
            new_people.append(p)
    
    return new_people, grid

# 可视化
fig, ax = plt.subplots()
im = ax.imshow(grid, cmap='jet', interpolation='nearest', vmin=0, vmax=1)

def animate(frame):
    global people, grid, im, field
    
    # 更新人群位置
    people, grid = update_people(grid, field, people)
    
    # 更新图像
    im.set_array(grid)
    
    # 清除并重新绘制
    ax.clear()
    ax.imshow(grid, cmap='jet', interpolation='nearest', vmin=0, vmax=1)
    
    # 绘制出口
    ax.scatter(*exit1, c='red', marker='s')
    ax.scatter(*exit2, c='red', marker='s')
    
    # 绘制轨迹（仅显示最新10步）
    for p in people:
        x, y = zip(*p['history'][-10:])  # 仅保留最近10步
        ax.plot(y, x, 'w-', alpha=0.5)
    
    # 显示疏散进度
    ax.text(0.05, 0.95, f'疏散进度: {len(people)} / {num_people}', transform=ax.transAxes, color='white')
    
    return [im]

ani = FuncAnimation(fig, animate, interval=100, blit=True)
plt.show()